import { randomUUID } from 'node:crypto';
import { IncomingMessage, Server } from 'node:http';
import { WebSocket, WebSocketServer } from 'ws';
const HEARTBEAT_INTERVAL = Number(process.env.HEARTBEAT_INTERVAL) || 28000;
interface BaseMessage<T, D> {
  id?: string;
  type: T;
  data: D;
}
export type ToolCallsMessage = BaseMessage<'toolCalls', ToolCall[]>;
export type TollCallMessage = BaseMessage<'toll_call', { name: string; arguments: any }>;
export type InitializeMess = BaseMessage<'initialize', { sessionId: string }>;
export type Message = InitializeMess | ToolCallsMessage;
export interface ToolCall {
  id: string;
  type: 'function';
  function?: FunctionCall;
}
interface FunctionCall {
  /** Function name */
  name: string;
  /** Function arguments */
  args: any;
}
// 日志函数
const logWs = (message: string, ...args: any[]) => {
  console.log(`[WS] ${message}`, ...args);
};

// 用于区分和管理多用户（每个连接分配一个唯一ID）
const clients = new Map<string, WebSocket>();
export const getSocket = (sessionId: string) => {
  return clients.get(sessionId);
};

export const registerWebSocket = (server: Server) => {
  const wss = new WebSocketServer({ server });
  wss.on('connection', handleConnection);
};
const handleConnection = (ws: WebSocket, req: IncomingMessage) => {
  const sessionId = randomUUID();
  clients.set(sessionId, ws);
  logWs(`客户端已连接，ID: ${sessionId}`);

  sendMessage(ws, {
    type: 'initialize',
    data: { sessionId },
  });

  // 心跳机制
  let isAlive = true;
  ws.on('pong', () => {
    isAlive = true;
  });
  const heartbeat = setInterval(() => {
    if (!isAlive) {
      logWs(`心跳超时，断开客户端: ${sessionId}`);
      ws.terminate();
      clearInterval(heartbeat);
      clients.delete(sessionId);
      return;
    }
    isAlive = false;
    ws.ping();
  }, HEARTBEAT_INTERVAL);

  ws.on('message', (message) => {
    logWs(`收到来自${sessionId}的消息:`, message.toString());
    // 回显消息并带上sessionId
    // ws.send(`服务端收到(${sessionId}): ${message}`);
  });

  ws.on('close', () => {
    clearInterval(heartbeat);
    clients.delete(sessionId);
    logWs(`客户端已断开连接，ID: ${sessionId}`);
  });
};
// 创建 WebSocket 服务器，绑定到 HTTP 服务器

export const sendMessage = (ws: WebSocket, message: Message) => {
  ws.send(JSON.stringify(message));
};
let nextId = 1;
export const call = (sessionId: string, name: string, args: any) => {
  const ws = getSocket(sessionId);
  if (!ws) throw new Error('No WebSocket connection found for session ID');
  const id = (nextId++).toString();
  sendMessage(ws, {
    type: 'toolCalls',
    data: [
      {
        id: id,
        type: 'function',
        function: { name, args },
      },
    ],
  });
  return { text: name, ok: true };
};
